from numpy import random as nrandom


class Task:
    def __init__(self, SliceNumber, LayerNumber, WorkLoad_mu, WorkLoad_sigma_square):
        self.LayerNumber = LayerNumber
        self.WorkloadList = list(map(abs, list(nrandom.normal(WorkLoad_mu, WorkLoad_sigma_square, LayerNumber))))
        self.WorkLoadSum = sum(self.WorkloadList)
        self.SliceNumber = SliceNumber
        self.Slice = [0 for _ in range(SliceNumber)]
        self.GetSliced()

    def SlicedWithMaxWorkLoad(self, MaxWorkLoad):
        self.Slice.clear()
        Sz = 0
        for i in range(len(self.WorkloadList)):
            if Sz + self.WorkloadList[i] <= MaxWorkLoad:
                Sz = self.WorkloadList[i] + Sz
            else:
                self.Slice.append(Sz)
                Sz = self.WorkloadList[i]
        if Sz != 0:
            self.Slice.append(Sz)

    def GetSliced(self):
        L = self.SliceNumber
        lp, rp = max(self.WorkloadList), sum(self.WorkloadList)
        while rp - lp > 0.00001:
            mid = (lp + rp) / 2
            self.SlicedWithMaxWorkLoad(mid)
            if len(self.Slice) > L:
                lp = mid
            else:
                rp = mid
        self.SlicedWithMaxWorkLoad(lp + 0.1)


class Satellite:
    def __init__(self, CPU_Number, SliceNumber, LayerNumber, MissionIncidence, WorkLoad_mu,
                 WorkLoad_sigma_square, MaxCPUWorkLoad, MaxAllowedWorkLoadofBlock):
        self.MaxCPUWorkLoad = MaxCPUWorkLoad
        self.CPU_WorkLoad_List = [0 for _ in range(CPU_Number)]
        self.TaskNumber = nrandom.poisson(MissionIncidence, 1)[0]
        self.TaskList = [Task(SliceNumber, LayerNumber, WorkLoad_mu, WorkLoad_sigma_square) for _ in
                         range(self.TaskNumber)]
        self.TaskBlock = []
        self.MaxAllowedWorkLoadofSlice = MaxAllowedWorkLoadofBlock
        self.CPU_Number = CPU_Number
        self.CPU_Loaded_list = [dict() for _ in range(self.CPU_Number)]
        self.DivideIntoBlocks()

    def ResidualResource(self):
        return 1 - sum(self.CPU_WorkLoad_List) / (self.CPU_Number * self.MaxCPUWorkLoad)

    def DivideIntoBlocks(self):
        cur, Sz = [], 0
        for (i, task) in enumerate(self.TaskList):
            if Sz + task.WorkLoadSum <= self.MaxAllowedWorkLoadofSlice:
                Sz = task.WorkLoadSum + Sz
                cur.append(i)
            else:
                self.TaskBlock.append(cur.copy())
                cur.clear()
                cur.append(i)
                Sz = task.WorkLoadSum
        if Sz != 0:
            self.TaskBlock.append(cur.copy())

    def SliceTasks(self):
        for block in self.TaskBlock:
            for task_idx in block:
                self.TaskList[task_idx].GetSliced()

    def CanLoad(self, cpu_idx, IncreasedWorkLoad):
        return IncreasedWorkLoad + self.CPU_WorkLoad_List[cpu_idx] <= self.MaxCPUWorkLoad

    def CPU_Load(self, cpu_idx, task_slice_idx, Workload):
        if task_slice_idx in self.CPU_Loaded_list[cpu_idx]:
            self.CPU_Loaded_list[cpu_idx][task_slice_idx] += Workload
        else:
            self.CPU_Loaded_list[cpu_idx][task_slice_idx] = Workload
        self.CPU_WorkLoad_List[cpu_idx] += Workload

    def CPU_offload(self, cpu_idx, task_slice_idx):
        if task_slice_idx in self.CPU_Loaded_list[cpu_idx]:
            self.CPU_WorkLoad_List[cpu_idx] -= self.CPU_Loaded_list[cpu_idx][task_slice_idx]
            self.CPU_Loaded_list[cpu_idx].pop(task_slice_idx)

    def Load(self, slice_tasks: list[(float, int)], AbandonSet: set):
        ProcessedTime, TotalDrop = 0, 0
        for workload, idx in slice_tasks:
            if idx in AbandonSet:
                continue
            else:
                find_a_cpu_to_load = False
                for i in range(self.CPU_Number):
                    if self.CanLoad(i, workload):
                        ProcessedTime += workload
                        self.CPU_Load(i, idx, workload)
                        find_a_cpu_to_load = True
                        break
                if not find_a_cpu_to_load:
                    TotalDrop += 1
                    AbandonSet.add(idx)
        # print('AA', TotalDrop, slice_tasks, AbandonSet)
        return ProcessedTime, TotalDrop

    def Offload(self, Offload_table: list[int]):
        for i in range(self.CPU_Number):
            for j in Offload_table:
                self.CPU_offload(i, j)
                # self.CPU_


class Network:
    def __init__(self, OrbitNumber, SatelliteNumber, CPU_Number, SliceNumber, LayerNumber,
                 MissionIncidence, BandWidth, WorkLoad_mu, WorkLoad_sigma_square, MaxCPUWorkLoad,
                 MaxAllowedWorkLoadofBlock, MaxCommunicationDistance):
        self.SliceNumber = SliceNumber
        self.OrbitNumber = OrbitNumber
        self.SatelliteNumber = SatelliteNumber
        self.TotalSatelliteNumber = OrbitNumber * SatelliteNumber
        self.MaxCommunicationDistance = MaxCommunicationDistance
        self.SatelliteGroup = [
            [Satellite(CPU_Number, SliceNumber, LayerNumber, MissionIncidence, WorkLoad_mu, WorkLoad_sigma_square,
                       MaxCPUWorkLoad, MaxAllowedWorkLoadofBlock) for _ in
             range(SatelliteNumber)] for _
            in
            range(OrbitNumber)]
        self.BandWithMatrix = [[0 for _ in range(self.TotalSatelliteNumber)] for _ in
                               range(self.TotalSatelliteNumber)]
        for i in range(self.TotalSatelliteNumber):
            for j in range(self.TotalSatelliteNumber):
                Rowi, Coli = i // self.SatelliteNumber, i % self.SatelliteNumber
                Rowj, Colj = j // self.SatelliteNumber, j % self.SatelliteNumber
                self.BandWithMatrix[i][j] = (abs(Rowj - Rowi) + abs(Colj - Coli)) * BandWidth

    def GetAdjSatellites(self, Row_idx, Col_idx):
        raw_res, D, N = [], self.MaxCommunicationDistance, self.OrbitNumber * self.SatelliteNumber
        for dr in range(1, D + 1):
            raw_res.extend([(Row_idx + dr, (Col_idx + dc + self.SatelliteNumber) % self.SatelliteNumber) for dc in
                            range(dr - D, D - dr + 1)])
            raw_res.extend([(Row_idx - dr, (Col_idx + dc + self.SatelliteNumber) % self.SatelliteNumber) for dc in
                            range(dr - D, D - dr + 1)])
        raw_res.extend(
            [(Row_idx, (self.SatelliteNumber + Col_idx + dc) % self.SatelliteNumber) for dc in range(-D, D + 1)])
        res = list(set([self.CoordinationMapping(r, c) for r, c in raw_res if N > self.CoordinationMapping(r, c) >= 0]))
        return res

    def Get(self, idx):
        x, y = self.Inv_CoordinationMapping(idx)
        return self.SatelliteGroup[x][y]

    def CoordinationMapping(self, Row_idx, Col_idx):
        return Row_idx * self.SatelliteNumber + Col_idx

    def Inv_CoordinationMapping(self, idx):
        return idx // self.SatelliteNumber, idx % self.SatelliteNumber

    def CalcDistance(self, idx1, idx2):
        row1, col1 = self.Inv_CoordinationMapping(idx1)
        row2, col2 = self.Inv_CoordinationMapping(idx2)
        return abs(row1 - row2) + abs(col1 + col2)

    def GetResource(self):
        S = sum([sum([satellite.ResidualResource() for satellite in orbit]) for orbit in self.SatelliteGroup])
        return S / self.TotalSatelliteNumber
# Z = Network(3, 4, 12, 4, 8, 10, 20, 10, 1, 500, 200)
# Around = Z.GetAdjSatellites(1, 1, 1)
# print(Around)
# for i in Around:
#     for j in Around:
#         print(i, j, Z.BandWithMatrix[i][j])
# x = Satellite(1, 4, 5, 20, 30, 10, 100, 600)
#
# for block in x.TaskBlock:
#     for task_index in block:
#         print(task_index,x.TaskList[task_index].Slice)
